#!/usr/bin/python
# 作者:龚华
# 2023年07月16日11时18分42秒
# -*- coding:utf-8 -*-
from math import log

import operator
import matplotlib.pyplot as plt
import pickle
import treePlotter
def calcShannonEnt(dataSet):
    numEntries = len(dataSet)
    labelCounts = {}

    #统计标签出现的次数
    #二维列表，可以用for循环遍历每一行的数据
    for featVec in dataSet:  # 每一行的数据，即每一个样本，featVec是一个列表
        currentLabel = featVec[-1]  # 取每一行的最后一个元素,即标签
        if currentLabel not in labelCounts.keys():
            labelCounts[currentLabel] = 0  # 如果这个元素不在labelCounts的key中，就把这个元素加入到labelCounts的key中
        labelCounts[currentLabel] += 1  # 如果这个元素在labelCounts的key中，就把这个元素的value加1
    shannonEnt = 0.0
    for key in labelCounts:
        prob = float(labelCounts[key]) / numEntries
        # print(prob)
        shannonEnt -= prob * log(prob, 2)  # log base 2
    return shannonEnt

def createDataSet():
    dataSet = [[1, 1, 'yes'],
               [1, 1, 'yes'],
               [1, 0, 'no' ],
               [0, 1, 'no' ],
               [0, 1, 'no' ]]
    labels = ['no surfacing', 'flippers']
    return dataSet, labels

def createDataSet2():
    dataSet = [[1, 1, 'maybe'],
               [1, 1, 'yes'],
               [1, 0, 'no' ],
               [0, 1, 'no' ],
               [0, 1, 'no' ]]
    labels = ['no surfacing', 'flippers']
    return dataSet, labels
#第三种创建数据集的方法。这个数据集是用于测试的
def retrieveTree(i):
    listOfTrees =[{'no surfacing': {0: 'no', 1: {'flippers': \
                    {0: 'no', 1: 'yes'}}}},
                  {'no surfacing': {0: 'no', 1: {'flippers': \
                    {0: {'head': {0: 'no', 1: 'yes'}}, 1: 'no'}}}}
                  ]
    return listOfTrees[i]
#分离
def splitDataSet(dataSet,axis,value):
    #创建list对象
    retDataSet=[]

    for featVec in dataSet:

        #分割符合特征的样本
        if featVec[axis] == value:  # 如果这个样本的第axis个特征等于value
            reducedFeatVec = featVec[:axis]  # 取featVec的前axis个特征
            reducedFeatVec.extend(featVec[axis+1:]) # 取featVec的后axis个特征
            #python list extend()函数用于在列表末尾一次性追加另一个序列中的多个值
            #python list append()函数用于在列表末尾添加新的对象,只能添加一个对象,这个对象可以是列表
            #[1,2,3].extend([4,5,6])=[1,2,3,4,5,6]
            #[1,2,3].append([4,5,6])=[1,2,3,[4,5,6]]
            retDataSet.append(reducedFeatVec)  # 把这个样本的前axis个特征和后axis个特征合并成一个样本，加入到retDataSet中
    return retDataSet
#将遍历整个数据集，循环计算香农熵和splitDataSet()函数，找到最好的特征划分方式
def chooseBestFeatureToSplit(dataSet):
    """
    计算信息增益，用ID3算法选择最好的数据集划分方式
    分割两个特征，比较哪个特征的信息增益最大，就用哪个特征来分割数据集
    """
    numFeatures = len(dataSet[0]) -1  #第一行的长度（既是列的个数）-1 =特征数
    baseEntropy = calcShannonEnt(dataSet)  #计算原始香农熵

    bestInfoGain = 0.0
    bestFeature = -1  #

    for i in range(numFeatures):

        #可以这样理解，dataSet是一个二维列表，每一行是一个样本，每一列是一个特征
        featList = [example[i] for example in dataSet] #取每一行的第i个元素，即第i个特征
        uniqueVals =set(featList)  #set()函数创建一个无序不重复元素集，用于去重，集合中的元素是唯一的
        newEntropy = 0.0

        #计算每种划分方式的信息熵；公式：信息熵=每种划分方式的信息熵*该划分方式的概率
        for value in uniqueVals:
            subDataSet = splitDataSet(dataSet,i,value)  #取出第i个特征等于value的样本

            prob = len(subDataSet) /float(len(dataSet))

            newEntropy += prob*calcShannonEnt(subDataSet)
                            #计算第i个特征等于value的样本的香农熵，乘以这个样本所占的比例，然后加到newEntropy中

        #第i个特征的信息增益
        infoGain =baseEntropy - newEntropy

        if (infoGain> bestInfoGain):
            bestInfoGain =infoGain
            bestFeature = i
    return bestFeature

def majorityCount(classList):
    """
    选择出现次数最多的一个结果
    投票表决，返回出现次数最多的分类名称

    """

    classCount ={}
    for vote in classList:
        if vote not in classCount.keys():  #如果这个分类不在classCount的key中，就把这个分类加入到classCount的key中
            classCount[vote]=0
        classCount[vote] +=1 #如果这个分类在classCount的key中，就把这个分类的value加1

    sortedClassCount = sorted(classCount.items(),key=operator.items(1),reverse=True)
        #classCount.items() 按值排序，返回的是一个列表，列表中的每个元素是一个元组，元组的第一个元素是key，第二个元素是value


    return sortedClassCount[0][0]  #返回出现次数最多的分类名称，返回最大的那个key

#创建树的函数代码
#用字典类型存储树的信息
def createTree(dataSet,labels):
    classList = [example[-1] for example in dataSet]  #取出所有标签

    #递归停止条件1：所有的类标签完全相同，则直接返回该类标签
    if classList.count(classList[0]) == len(classList):  #如果所有的分类都是一样的，就返回这个分类
        return classList[0]
    #递归停止条件2：使用完了所有特征，仍然不能将数据集划分成仅包含唯一类别的分组
    if len(dataSet[0]) == 1:  #如果只有一个特征，就返回出现次数最多的分类
        return majorityCount(classList)  #返回出现次数最多的分类名称，调用函数majorityCount()


    bestFeat = chooseBestFeatureToSplit(dataSet)  #选择最好的特征
    bestFeatLabel = labels[bestFeat]  #最好的特征的名称

    #字典类型，存储树的信息
    myTree = {bestFeatLabel:{}}  #创建一个字典，字典的key是最好的特征的名称，value是一个字典


    del(labels[bestFeat])  #删除labels中的最好的特征
    featValues = [example[bestFeat] for example in dataSet]  #取出最好的特征的所有值
    uniqueVals = set(featValues)  #去重

    for value in uniqueVals:  #遍历最好的特征的所有值
        subLabels = labels[:]  #复制labels，因为python函数参数是列表类型时，参数是按照引用方式传递的，为了保护原始列表，复制一份

       #递归createTree()函数，为myTree添加分支节点
        myTree[bestFeatLabel][value] = createTree(splitDataSet(dataSet,bestFeat,value),subLabels)
            #递归创建树，每次递归调用时，传入的参数是最好的特征的所有值，和复制的labels
    return myTree


#使用决策树的分类函数
def classify(inputTree,featLabels,testVec):
    #inputTree是已经生成的决策树，featLabels是标签列表，testVec是测试数据列表
    #决策树形式：{'no surfacing': {0: 'no', 1: {'flippers': {0: 'no', 1: 'yes'}}}}

    firstStr = list(inputTree.keys())[0]  #取出第一个key，即第一个特征
    secondDict = inputTree[firstStr]  #取出第一个key的value，即第一个特征的所有值,是一个字典

    featIndex = featLabels.index(firstStr)  #第一个特征的索引
    for key in secondDict.keys():
        if testVec[featIndex] == key:
            if type(secondDict[key]).__name__ =='dict':
                #如果secondDict[key]是字典类型，就递归调用classify()函数
                classLabel = classify(secondDict[key],featLabels,testVec)
            else:
                classLabel = secondDict[key]
    return classLabel


#使用pickle模块存储决策树,
#决策树建立过程很耗时，使用pickle模块存储决策树，使用时直接调用即可
def storeTree(inputTree,filename):
    #import pickle
    fw = open(filename,'wb')
    pickle.dump(inputTree,fw)
    fw.close()

def grabTree(filename):
    #import pickle
    fr = open(filename,'rb')
    return pickle.load(fr)


#使用决策树预测隐形眼镜类型
def lensesTest():
    fr=open('.\data\lenses.txt')
    lenses=[inst.strip().split('\t') for inst in fr.readlines()]
    lensesLabels=['age','prescript','astigmatic','tearRate']
    lensesTree = createTree(lenses,lensesLabels)
    print(lensesTree)
    treePlotter.createPlot(lensesTree)




#测试函数部分
def test_calcShannonEnt():
    myDat, labels = createDataSet()
    print(calcShannonEnt(myDat))
    print(labels)
    #熵越高，混合的数据也越多，我们可以在数据集中添加更多的分类，观察熵是如何变化的
    maDat, labels = createDataSet2()
    print(calcShannonEnt(maDat))
    print(labels)

def test_splitDataSet():
    myDat, labels = createDataSet()
    print(myDat)
    print(splitDataSet(myDat,0,1))
    print(splitDataSet(myDat,0,0))
def test_chooseBestFeatureToSplit():
    myDat, labels = createDataSet()
    print(chooseBestFeatureToSplit(myDat))
    print(labels)
    print(myDat)
def test_createTree():
    myDat, labels = createDataSet()
    myTree = createTree(myDat,labels)
    print(myTree)

def test_classify():
    myDat,labels = createDataSet()
    print(labels)
    print(myDat)
    myTree = retrieveTree(0)
    print(myTree)
    print(classify(myTree,labels,[1,0]))
    print(classify(myTree,labels,[1,1]))

#测试决策树的存储
def test_storeTree():
    myTree = retrieveTree(0)
    print(myTree)
    storeTree(myTree,'.\data\classifierStorage.txt')
    print(grabTree('.\data\classifierStorage.txt'))
def test_grabTree():
    print(grabTree('.\data\classifierStorage.txt'))

if __name__ == '__main__':
    #测试信息熵
    #test_calcShannonEnt()
    #测试划分数据集
    #test_splitDataSet()

    #测试选择最好的数据集划分方式
    #test_chooseBestFeatureToSplit()

    #测试创建树
    #test_createTree()

    #测试使用决策树的分类函数
    #test_classify()

    #测试决策树的存储
    #test_storeTree()
    #test_grabTree()

    #使用决策树预测隐形眼镜类型
    lensesTest()